#define STB_IMAGE_IMPLEMENTATION
#include "../stb_image.h"
#include "BufferManager.h"
#include <iostream>
#include <fstream>
#include <cstring>

BufferManager::BufferManager(VulkanApplication* app)
    : _descriptorPool(nullptr), _descriptorSetLayout(nullptr), _application(app)
{}

BufferManager::~BufferManager()
{}

bool BufferManager::prepare(const std::string& vertFile, const std::string& fragFile)
{
    VkVertexInputBindingDescription bindingDescription
    {
        0, sizeof(Vertex),
        VK_VERTEX_INPUT_RATE_VERTEX
    };

    std::vector<VkVertexInputAttributeDescription> attrDescription;
    attrDescription.push_back(
    {
        0, 0,  // location, binding
        VK_FORMAT_R32G32B32_SFLOAT,  // format
        0                         // offset
    });
    attrDescription.push_back(
    {
        1, 0,  // location, binding
        VK_FORMAT_R32G32B32_SFLOAT,  // format
        sizeof(float) * 3         // offset
    });
    attrDescription.push_back(
    {
        2, 0,  // location, binding
        VK_FORMAT_R32G32_SFLOAT,  // format
        sizeof(float) * 6         // offset
    });

    VkPipelineVertexInputStateCreateInfo inputVertexInfo
    {
        VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO,
        nullptr, 0,
        1, &bindingDescription,
        (uint32_t)attrDescription.size(), attrDescription.data()
    };

    VkPipelineInputAssemblyStateCreateInfo inputAssemblyInfo
    {
        VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO,
        nullptr, 0,
        VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST, VK_FALSE
    };

    // Create graphics pipeline
    return _application->setUpPipeline(
        vertFile, fragFile, inputVertexInfo, inputAssemblyInfo, _descriptorSetLayout);
}

bool BufferManager::createDescriptorSet(int maxSets)
{
    // Create uniform desciption layout
    VkDescriptorSetLayoutBinding uboLayoutBinding
    {
        0,  // binding
        VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
        1,  // descriptorCount
        VK_SHADER_STAGE_VERTEX_BIT,
        nullptr
    };

    VkDescriptorSetLayoutBinding samplerLayoutBinding
    {
        1,  // binding
        VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
        1,  // descriptorCount,
        VK_SHADER_STAGE_FRAGMENT_BIT,
        nullptr
    };

    VkDescriptorSetLayoutBinding bindings[2] = { uboLayoutBinding, samplerLayoutBinding };
    VkDescriptorSetLayoutCreateInfo layoutInfo
    {
        VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO,
        nullptr, 0,
        2, bindings
    };

    VkResult result = vkCreateDescriptorSetLayout(
        _application->getDevice(), &layoutInfo, nullptr, &_descriptorSetLayout);
    if (result != VK_SUCCESS)
    {
        printf("Failed to create description-set layout.\n");
        return false;
    }

    // Create uniform desciption pool
    VkDescriptorPoolSize poolSizes[2] =
    {
        { VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 1 },
        { VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, 1 }
    };

    VkDescriptorPoolCreateInfo poolInfo
    {
        VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO,
        nullptr, 0,
        (unsigned int)maxSets, 2, poolSizes
    };

    result = vkCreateDescriptorPool(
        _application->getDevice(), &poolInfo, nullptr, &_descriptorPool);
    if (result != VK_SUCCESS)
    {
        printf("Failed to create description pool.\n");
        return false;
    }
    return true;
}

bool BufferManager::allocateNewDescriptor(const std::string& name)
{
    // Create uniform description sets
    VkDescriptorSetAllocateInfo allocInfo
    {
        VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO,
        nullptr,
        _descriptorPool,
        1, &_descriptorSetLayout
    };

    VkDescriptorSet descriptorSet = nullptr;
    VkResult result = vkAllocateDescriptorSets(_application->getDevice(), &allocInfo, &descriptorSet);
    if (result != VK_SUCCESS)
    {
        printf("Failed to create description set.\n");
        return false;
    }
    _descriptorSetMap[name] = descriptorSet;
    return true;
}

void BufferManager::cleanup()
{
    for (std::map<std::string, VkBuffer>::iterator itr = _bufferMap.begin();
        itr != _bufferMap.end(); ++itr)
    {
        vkDestroyBuffer(_application->getDevice(), itr->second, nullptr);
    }

    for (std::map<std::string, VkDeviceMemory>::iterator itr = _bufferMemoryMap.begin();
         itr != _bufferMemoryMap.end(); ++itr)
    {
        vkFreeMemory(_application->getDevice(), itr->second, nullptr);
    }

    for (std::map<std::string, VkSampler>::iterator itr = _imageSamplerMap.begin();
         itr != _imageSamplerMap.end(); ++itr)
    {
        vkDestroySampler(_application->getDevice(), itr->second, nullptr);
    }

    for (std::map<std::string, VkImageView>::iterator itr = _imageViewMap.begin();
         itr != _imageViewMap.end(); ++itr)
    {
        vkDestroyImageView(_application->getDevice(), itr->second, nullptr);
    }
    
    for (std::map<std::string, VkImage>::iterator itr = _imageMap.begin();
         itr != _imageMap.end(); ++itr)
    {
        vkDestroyImage(_application->getDevice(), itr->second, nullptr);
    }
    vkDestroyDescriptorSetLayout(_application->getDevice(), _descriptorSetLayout, nullptr);
    vkDestroyDescriptorPool(_application->getDevice(), _descriptorPool, nullptr);
}

VkBuffer BufferManager::createVertexBuffer(const std::string& name, const std::vector<Vertex>& vertices)
{
    size_t dataSize = sizeof(Vertex) * vertices.size();
    VkBuffer vertexBuffer = nullptr;
    VkDeviceMemory vertexBufferMemory = nullptr;
    if (!createBuffer(vertexBuffer, dataSize, VK_BUFFER_USAGE_VERTEX_BUFFER_BIT))
        return nullptr;

    VkMemoryRequirements memRequirements;
    vkGetBufferMemoryRequirements(_application->getDevice(), vertexBuffer, &memRequirements);
    if (!createBufferMemory(vertexBufferMemory, memRequirements,
                            VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT))
        return nullptr;

    vkBindBufferMemory(_application->getDevice(), vertexBuffer, vertexBufferMemory, 0);
    _bufferMap[name] = vertexBuffer;
    _bufferMemoryMap[name] = vertexBufferMemory;

    void* data = nullptr;
    vkMapMemory(_application->getDevice(), vertexBufferMemory, 0, dataSize, 0, &data);
    memcpy(data, vertices.data(), dataSize);
    vkUnmapMemory(_application->getDevice(), vertexBufferMemory);
    return vertexBuffer;
}

VkBuffer BufferManager::createIndexBuffer(const std::string& name, const std::vector<unsigned int>& indices)
{
    size_t dataSize = sizeof(unsigned int) * indices.size();
    VkBuffer indexBuffer = nullptr;
    VkDeviceMemory indexBufferMemory = nullptr;
    if (!createBuffer(indexBuffer, dataSize, VK_BUFFER_USAGE_INDEX_BUFFER_BIT))
        return nullptr;

    VkMemoryRequirements memRequirements;
    vkGetBufferMemoryRequirements(_application->getDevice(), indexBuffer, &memRequirements);
    if (!createBufferMemory(indexBufferMemory, memRequirements,
                            VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT))
        return nullptr;

    vkBindBufferMemory(_application->getDevice(), indexBuffer, indexBufferMemory, 0);
    _bufferMap[name] = indexBuffer;
    _bufferMemoryMap[name] = indexBufferMemory;

    void* data = nullptr;
    vkMapMemory(_application->getDevice(), indexBufferMemory, 0, dataSize, 0, &data);
    memcpy(data, indices.data(), dataSize);
    vkUnmapMemory(_application->getDevice(), indexBufferMemory);
    return indexBuffer;
}

VkBuffer BufferManager::createUniformBuffer(const std::string& name, size_t uboSize)
{
    VkBuffer uniformBuffer = nullptr;
    VkDeviceMemory uniformBufferMemory = nullptr;
    if (!createBuffer(uniformBuffer, uboSize, VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT))
        return nullptr;

    VkMemoryRequirements memRequirements;
    vkGetBufferMemoryRequirements(_application->getDevice(), uniformBuffer, &memRequirements);
    if (!createBufferMemory(uniformBufferMemory, memRequirements,
                            VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT))
        return nullptr;
    
    vkBindBufferMemory(_application->getDevice(), uniformBuffer, uniformBufferMemory, 0);
    _bufferMap[name] = uniformBuffer;
    _bufferMemoryMap[name] = uniformBufferMemory;
    return uniformBuffer;
}

void BufferManager::updateUniformDescriptor(const std::string& desc, const std::string& uniform, size_t uboSize)
{
    VkBuffer uniformBuffer = _bufferMap[uniform];
    VkDescriptorBufferInfo bufferInfo
    { uniformBuffer, 0, uboSize };

    VkWriteDescriptorSet descriptorWrite
    {
        VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
        nullptr,
        _descriptorSetMap[desc], 0,  // dstSet, dstBinding
        0, 1,                        // dstArrayElement, descriptorCount
        VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
        nullptr, &bufferInfo, nullptr
    };
    vkUpdateDescriptorSets(_application->getDevice(), 1, &descriptorWrite, 0, nullptr);
}

VkImage BufferManager::createTextureImage(const std::string& name, const std::string& filename)
{
    int w = 0, h = 0, ch = 4;
    stbi_uc* pixels = stbi_load(filename.c_str(), &w, &h, &ch, STBI_rgb_alpha);
    if (!pixels)
    {
        printf("Failed to load image file: %s\n", filename.c_str());
        return nullptr;
    }

    VkImage image = createImage(name, pixels, w, h, 4);
    stbi_image_free(pixels);
    return image;
}

VkImage BufferManager::createTextureImage(const std::string& name, unsigned char* buffer, int size)
{
    int w = 0, h = 0, ch = 4;
    stbi_uc* pixels = stbi_load_from_memory(buffer, size, &w, &h, &ch, STBI_rgb_alpha);
    if (!pixels)
    {
        printf("Failed to load image memory.\n");
        return nullptr;
    }

    VkImage image = createImage(name, pixels, w, h, 4);
    stbi_image_free(pixels);
    return image;
}

VkImage BufferManager::createImage(const std::string& name, unsigned char* pixels,
                                   int w, int h, int ch)
{
    VkFormat format = VK_FORMAT_R8G8B8A8_SRGB;
    if (ch == 3) format = VK_FORMAT_R8G8B8_SRGB;
    
    // Create image buffer
    VkDeviceSize imageSize = w * h * ch;
    VkBuffer imageBuffer = nullptr;
    VkDeviceMemory imageBufferMemory = nullptr;
    if (!createBuffer(imageBuffer, imageSize, VK_BUFFER_USAGE_TRANSFER_SRC_BIT))
        return nullptr;

    VkMemoryRequirements memRequirements;
    vkGetBufferMemoryRequirements(_application->getDevice(), imageBuffer, &memRequirements);
    if (!createBufferMemory(imageBufferMemory, memRequirements,
                            VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT))
        return nullptr;

    void* imageData = nullptr;
    vkBindBufferMemory(_application->getDevice(), imageBuffer, imageBufferMemory, 0);
    vkMapMemory(_application->getDevice(), imageBufferMemory, 0, imageSize, 0, &imageData);
    memcpy(imageData, pixels, static_cast<size_t>(imageSize));
    vkUnmapMemory(_application->getDevice(), imageBufferMemory);

    // Create image object
    VkImageCreateInfo imageInfo
    {
        VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO,
        nullptr, 0,
        VK_IMAGE_TYPE_2D, format,
        { (uint32_t)w, (uint32_t)h, 1 },
        1, 1,    // mipLevels, arrayLayers
        VK_SAMPLE_COUNT_1_BIT, VK_IMAGE_TILING_OPTIMAL,
        VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_SAMPLED_BIT,
        VK_SHARING_MODE_EXCLUSIVE,
        0, nullptr,    // queueFamilyIndices
        VK_IMAGE_LAYOUT_UNDEFINED
    };

    VkImage image = nullptr;
    VkDeviceMemory imageResultBufferMemory = nullptr;
    VkResult result = vkCreateImage(_application->getDevice(), &imageInfo, nullptr, &image);
    if (result != VK_SUCCESS)
    {
        printf("Failed to create image object.\n");
        return nullptr;
    }

    VkMemoryRequirements memRequirements2;
    vkGetImageMemoryRequirements(_application->getDevice(), image, &memRequirements2);
    if (!createBufferMemory(imageResultBufferMemory, memRequirements2,
                            VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT))
        return nullptr;
    else
        vkBindImageMemory(_application->getDevice(), image, imageResultBufferMemory, 0);
    
    // Start copying command
    VkCommandBufferAllocateInfo allocInfo
    {
        VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO,
        nullptr, _application->getCommandPool(),
        VK_COMMAND_BUFFER_LEVEL_PRIMARY, 1
    };
    VkCommandBuffer tempCommand;
    vkAllocateCommandBuffers(_application->getDevice(), &allocInfo, &tempCommand);

    VkCommandBufferBeginInfo beginInfo
    {
        VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO,
        nullptr, VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT, nullptr
    };
    vkBeginCommandBuffer(tempCommand, &beginInfo);

    // Set image layout to transfer-destination
    VkImageMemoryBarrier barrierPre
    {
        VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER,
        nullptr, 0, VK_ACCESS_TRANSFER_WRITE_BIT,
        VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
        VK_QUEUE_FAMILY_IGNORED, VK_QUEUE_FAMILY_IGNORED,
        image, { VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1 }
    };
    vkCmdPipelineBarrier(tempCommand, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT,
                         0, 0, nullptr, 0, nullptr, 1, &barrierPre);

    // Copy from staging buffer to image result buffer
    VkBufferImageCopy region
    {
        0, 0, 0,
        { VK_IMAGE_ASPECT_COLOR_BIT, 0, 0, 1 },
        { 0, 0, 0 },
        { (uint32_t)w, (uint32_t)h, 1 }
    };
    vkCmdCopyBufferToImage(tempCommand, imageBuffer, image,
                           VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, &region);

    // Set image layout to sampling-ready
    VkImageMemoryBarrier barrierPost
    {
        VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER,
        nullptr, VK_ACCESS_TRANSFER_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT,
        VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
        VK_QUEUE_FAMILY_IGNORED, VK_QUEUE_FAMILY_IGNORED,
        image, { VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1 }
    };
    vkCmdPipelineBarrier(tempCommand, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT,
                         0, 0, nullptr, 0, nullptr, 1, &barrierPost);

    // Finish copying command and release resources
    VkQueue queue = nullptr;
    vkGetDeviceQueue(_application->getDevice(), 0, 0, &queue);
    vkEndCommandBuffer(tempCommand);

    VkSubmitInfo submitInfo
    {
        VK_STRUCTURE_TYPE_SUBMIT_INFO,
        nullptr,
        0, nullptr, nullptr,  // waitSemaphores
        1, &tempCommand,
        0, nullptr    // signalSemaphore
    };
    vkQueueSubmit(queue, 1, &submitInfo, VK_NULL_HANDLE);
    vkQueueWaitIdle(queue);

    vkFreeCommandBuffers(_application->getDevice(), _application->getCommandPool(), 1, &tempCommand);
    vkFreeMemory(_application->getDevice(), imageBufferMemory, nullptr);
    vkDestroyBuffer(_application->getDevice(), imageBuffer, nullptr);
    
    _imageMap[name] = image;
    _bufferMemoryMap[name] = imageResultBufferMemory;
    return image;
}

VkSampler BufferManager::createTextureSampler(const std::string& name, VkImage image, VkFormat format)
{
    VkImageViewCreateInfo viewInfo
    {
        VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO,
        nullptr, 0,
        image, VK_IMAGE_VIEW_TYPE_2D,
        format,
        { VK_COMPONENT_SWIZZLE_IDENTITY, VK_COMPONENT_SWIZZLE_IDENTITY,
          VK_COMPONENT_SWIZZLE_IDENTITY, VK_COMPONENT_SWIZZLE_IDENTITY },
        { VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1 }
    };

    VkImageView imageView = nullptr;
    VkResult result = vkCreateImageView(_application->getDevice(), &viewInfo, nullptr, &imageView);
    if (result != VK_SUCCESS)
    {
        printf("Failed to create texture image view.\n");
        return nullptr;
    }

    VkPhysicalDeviceProperties properties;
    vkGetPhysicalDeviceProperties(_application->getPhysicalDevice(), &properties);

    VkSamplerCreateInfo samplerInfo
    {
        VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO,
        nullptr, 0,
        VK_FILTER_LINEAR, VK_FILTER_LINEAR,
        VK_SAMPLER_MIPMAP_MODE_LINEAR,
        VK_SAMPLER_ADDRESS_MODE_REPEAT,
        VK_SAMPLER_ADDRESS_MODE_REPEAT,
        VK_SAMPLER_ADDRESS_MODE_REPEAT,
        0.0f,  // lodBias
        VK_TRUE, properties.limits.maxSamplerAnisotropy,  // anisotropy
        VK_FALSE, VK_COMPARE_OP_ALWAYS,  // compare op
        0.0f, 0.0f,  // minLod, maxLod
        VK_BORDER_COLOR_INT_OPAQUE_BLACK,
        VK_FALSE  // unnormalized coordinates
    };

    VkSampler texSampler = nullptr;
    result = vkCreateSampler(_application->getDevice(), &samplerInfo, nullptr, &texSampler);
    if (result != VK_SUCCESS)
    {
        printf("Failed to create texture sampler.\n");
        return nullptr;
    }

    _imageViewMap[name] = imageView;
    _imageSamplerMap[name] = texSampler;
    return texSampler;
}

void BufferManager::updateTextureDescriptor(const std::string& desc, const std::string& texture,
                                            VkImageLayout layout)
{
    VkImageView textureImageView = _imageViewMap[texture];
    VkSampler textureSampler = _imageSamplerMap[texture];
    VkDescriptorImageInfo imageInfo
    {
        textureSampler, textureImageView, layout
    };

    VkWriteDescriptorSet descriptorWrite
    {
        VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
        nullptr,
        _descriptorSetMap[desc], 1,  // dstSet, dstBinding
        0, 1,                        // dstArrayElement, descriptorCount
        VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
        &imageInfo, nullptr, nullptr
    };
    vkUpdateDescriptorSets(_application->getDevice(), 1, &descriptorWrite, 0, nullptr);
}

void* BufferManager::mapBufferMemory(const std::string& name, size_t dataSize)
{
    void* dataPtr = nullptr;
    if (_bufferMemoryMap.find(name) != _bufferMemoryMap.end())
    {
        vkMapMemory(_application->getDevice(), _bufferMemoryMap[name],
                    0, dataSize, 0, &dataPtr);
    }
    return dataPtr;
}

void BufferManager::unmapBufferMemory(const std::string& name)
{
    if (_bufferMemoryMap.find(name) != _bufferMemoryMap.end())
        vkUnmapMemory(_application->getDevice(), _bufferMemoryMap[name]);
}

VkBuffer BufferManager::getBuffer(const std::string& name)
{
    std::map<std::string, VkBuffer>::iterator itr = _bufferMap.find(name);
    if (itr != _bufferMap.end()) return itr->second;
    else return nullptr;
}

bool BufferManager::createBuffer(VkBuffer& buffer, size_t dataSize, VkBufferUsageFlags usage)
{
    VkBufferCreateInfo bufferInfo
    {
        VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO,
        nullptr, 0,
        dataSize,
        usage,
        VK_SHARING_MODE_EXCLUSIVE,
        0, nullptr  // pQueueFamilyIndices
    };

    VkResult result = vkCreateBuffer(_application->getDevice(), &bufferInfo, nullptr, &buffer);
    if (result != VK_SUCCESS)
    {
        printf("Failed to create data buffer.\n");
        return false;
    }
    return true;
}

bool BufferManager::createBufferMemory(VkDeviceMemory& bufferMemory,
                                       VkMemoryRequirements& memRequirements, uint32_t propertyBits)
{
    VkPhysicalDeviceMemoryProperties memProperties;
    vkGetPhysicalDeviceMemoryProperties(_application->getPhysicalDevice(), &memProperties);

    uint32_t memoryTypeIndex = 0;
    for (uint32_t i = 0; i < memProperties.memoryTypeCount; ++i)
    {
        if ((memRequirements.memoryTypeBits & (1 << i)) &&
            (memProperties.memoryTypes[i].propertyFlags & propertyBits) == propertyBits)
        { memoryTypeIndex = i; }
    }

    VkMemoryAllocateInfo allocInfo
    {
        VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO,
        nullptr,
        memRequirements.size,
        memoryTypeIndex
    };

    VkResult result = vkAllocateMemory(_application->getDevice(), &allocInfo, nullptr, &bufferMemory);
    if (result != VK_SUCCESS)
    {
        printf("Failed to allocate buffer memory.\n");
        return false;
    }
    return true;
}
